import cv2
import numpy as np
import DirectionControl
import matplotlib.pyplot as plt
# Step1. 转换为HSV
import time


def fuction(pic, flag):
    # real_image为真实图片
    pic = cv2.resize(pic, (640, 480), interpolation=cv2.INTER_CUBIC)
    real_image = pic
# temp_image为在中间步骤进行处理的图片
    temp_image = pic
    temp_image = cv2.cvtColor(temp_image, cv2.COLOR_BGR2HSV)

    # 黄色
    low_range_yellow = np.array([20, 50, 40])
    high_range_yellow = np.array([40, 150, 200])

    # 绿色
    low_range_green = np.array([70, 40, 40])
    high_range_green = np.array([85, 100, 200])

    # 红色
    low_range_red = np.array([170, 20, 40])
    high_range_red = np.array([180, 130, 200])
    # mask遮罩
    mask_yellow = cv2.inRange(temp_image, low_range_yellow, high_range_yellow)
    mask_green = cv2.inRange(temp_image, low_range_green, high_range_green)
    mask_red = cv2.inRange(temp_image, low_range_red, high_range_red)
    # cv2.imshow('Mask', mask_green)
    # cv2.waitKey(0)

    # 滤波去除噪声
    i = 0
    while i < 50:
        mask_yellow = cv2.medianBlur(mask_yellow, 5)
        mask_green = cv2.medianBlur(mask_green, 5)
        mask_red = cv2.medianBlur(mask_red, 5)
        i += 1

    # 进行一次膨胀操作
    mask_yellow = cv2.dilate(
        mask_yellow, cv2.getStructuringElement(
            cv2.MORPH_ELLIPSE, (3, 3)), iterations=5)
    mask_green = cv2.dilate(
        mask_green, cv2.getStructuringElement(
            cv2.MORPH_ELLIPSE, (3, 3)), iterations=5)
    mask_red = cv2.dilate(
        mask_red, cv2.getStructuringElement(
            cv2.MORPH_ELLIPSE, (3, 3)), iterations=5)

    # data_yellow=[]
    # data_green=[]
    # data_red=[]
    # data_temp=[]

    def draw_min_rectangle():  # conts = contours

        contours_yellow, hierarchy_yellow = cv2.findContours(
            mask_yellow, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        contours_green, hierarchy_green = cv2.findContours(
            mask_green, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        contours_red, hierarchy_red = cv2.findContours(
            mask_red, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        # 识别黄色
        data_temp = [320, 240, 0, 0]
        if flag == '1':
            for cnt in contours_yellow:
                x, y, w, h = cv2.boundingRect(cnt)
                if (w + h) < 50:
                    continue
                else:
                    # print(flag)
                    data_temp = [x, y, w, h]
                    # data_yellow.append(data_temp)
                    cv2.rectangle(
                        real_image, (x, y), (x + w, y + h), (125, 214, 242), 2)  # blue
            return data_temp
        # 识别绿色
        if flag == '2':
            for cnt in contours_green:
                x, y, w, h = cv2.boundingRect(cnt)
                if (w + h) < 70:
                    continue
                else:
                    if(flag == 2):
                        data_temp = [x, y, w, h]
                        # data_green.append(data_temp)
                        cv2.rectangle(
                            real_image, (x, y), (x + w, y + h), (0, 255, 0), 2)  # blue
            return data_temp
        # 识别红色
        if flag == '3':
            for cnt in contours_red:
                x, y, w, h = cv2.boundingRect(cnt)
                if (w + h) < 50:
                    continue
                else:
                    if(flag == 3):
                        data_temp = [x, y, w, h]
                        # data_red.append(data_temp)
                        cv2.rectangle(
                            real_image, (x, y), (x + w, y + h), (0, 0, 255), 2)  # blue
            return data_temp
    data = draw_min_rectangle()
    # draw_min_rectangle()
    print(data)
    cv2.imshow('Mask', real_image)
    return data

    # cv2.destroyAllWindows()
    # 输出区域信息
# img = cv2.imread('./pic/50.0.jpg')
# fuction(img,'1')
